# -*- coding: utf-8 -*-

import os
import tensorflow as tf


def get_images_path(folder):
    sub_folders = [os.path.join(folder, sub_folder) for sub_folder in os.listdir(folder)
                   if os.path.isdir(os.path.join(folder, sub_folder))]
    full_images_path, full_mask_images_path = [], []
    for sub_folder in sub_folders:
        images_path = [os.path.join(sub_folder, name) for name in os.listdir(sub_folder)
                       if "jpg" in name and "mask" not in name]
        mask_images_path = [path.replace(".jpg", "_mask.jpg") for path in images_path]
        full_images_path.extend(images_path)
        full_mask_images_path.extend(mask_images_path)
    assert len(full_images_path) == len(full_mask_images_path)
    return full_images_path, full_mask_images_path


def process_image(image_path, mask_image_path):
    image = tf.image.decode_jpeg(tf.read_file(image_path), channels=3)  # use rgb format
    image = tf.div(tf.sub(tf.to_float(image), [104, 110, 109]), 255)
    mask_image = tf.image.decode_jpeg(tf.read_file(mask_image_path), channels=1)
    mask_image = tf.to_float(mask_image)
    return image, mask_image


def generate_image_batch(folder, num_epochs=1, batch_size=32):
    images_path, mask_images_path = get_images_path(folder)
    image_path, mask_image_path = tf.train.slice_input_producer(
        [tf.convert_to_tensor(images_path, tf.string), tf.convert_to_tensor(mask_images_path, tf.string)],
        num_epochs=num_epochs)
    processed_image, processed_mask_image = process_image(image_path, mask_image_path)
    processed_mask_image = tf.image.resize_images(processed_mask_image, new_height=45, new_width=80)
    image_batch, mask_image_batch = tf.train.shuffle_batch([processed_image, processed_mask_image],
                                                           batch_size=batch_size,
                                                           capacity=batch_size * 10,
                                                           min_after_dequeue=batch_size * 5,
                                                           num_threads=4,
                                                           shapes=[[180, 320, 3], [45, 80, 1]])
    return image_batch, mask_image_batch


def main():
    folder = "/home/lijun/Dataset/vehicle_segment/train"
    image_batch_op, mask_image_batch_op = generate_image_batch(folder)
    coord = tf.train.Coordinator()
    with tf.Session() as session:
        tf.initialize_all_variables().run()
        tf.initialize_local_variables().run()
        threads = tf.train.start_queue_runners(sess=session, coord=coord)
        try:
            while not coord.should_stop():
                images, mask_images = session.run([image_batch_op, mask_image_batch_op])
                print(images.shape, mask_images.shape)
                print(images.dtype, mask_images.dtype)
                break
        except tf.errors.OutOfRangeError:
            print("all images have been used")
        finally:
            coord.request_stop()
        coord.join(threads)

if __name__ == "__main__":
    main()
